package org.flint.data.structure.leetcode;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import java.util.List;
import java.util.PriorityQueue;
import java.util.TreeMap;

/**
 * @author flint92
 */
public class TopFrequentElements {

    /**
     * Given a non-empty array of integers, return the k most frequent elements.
     * <p>
     * Example 1:
     * <p>
     * Input: nums = [1,1,1,2,2,3], k = 2
     * Output: [1,2]
     * Example 2:
     * <p>
     * Input: nums = [1], k = 1
     * Output: [1]
     * Note:
     * <p>
     * You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
     * Your algorithm's time complexity must be better than O(n log n), where n is the array's size.
     *
     * @param nums 数组
     * @param k    前k个
     * @return 集合
     */
    public List<Integer> topKFrequent(int[] nums, int k) {
        if (nums == null || nums.length == 0 || k <= 0) {
            return ImmutableList.of();
        }

        final TreeMap<Integer, Integer> map = buildTreeMap(nums);

        final PriorityQueue<Freq> queue = new PriorityQueue<>();
        for (Integer key : map.keySet()) {
            if (queue.size() < k) {
                queue.offer(new Freq(key, map.get(key)));
            } else {
                Freq front = queue.peek();
                if (map.get(key) >= front.freq) {
                    queue.poll();
                    queue.offer(new Freq(key, map.get(key)));
                }
            }
        }

        List<Integer> list = Lists.newLinkedList();
        while (!queue.isEmpty()) {
            list.add(queue.poll().e);
        }
        return list;
    }

    private TreeMap<Integer, Integer> buildTreeMap(int[] nums) {
        TreeMap<Integer, Integer> map = Maps.newTreeMap();
        for (int num : nums) {
            if (map.containsKey(num)) {
                map.put(num, map.get(num) + 1);
            } else {
                map.put(num, 1);
            }
        }
        return map;
    }

    private class Freq implements Comparable<Freq> {
        int e, freq;

        Freq(int e, int freq) {
            this.e = e;
            this.freq = freq;
        }

        @Override
        public int compareTo(Freq another) {
            return Integer.compare(freq, another.freq);
        }
    }

}
